from setting import backtest
from setting import data
import polars as pl
import datetime
import pickle
from setting.crate import *

# 策略名称 - 用于生成测试报告
strategy_name = "Alice22"

# 合约配置
base_data_symbol = ["SHFE_ag00_60"]  # 基准合约列表(用于时间轴)

trade_data_symbol = [
    "CZCE_AP00_60", "CZCE_CF00_60", "CZCE_CJ00_60", "CZCE_CY00_60", "CZCE_FG00_60",
    "CZCE_LR00_60", "CZCE_MA00_60", "CZCE_OI00_60", "CZCE_PF00_60", "CZCE_PK00_60",
    "CZCE_PL00_60", "CZCE_PM00_60", "CZCE_PR00_60", "CZCE_RI00_60", "CZCE_RM00_60",
    "CZCE_RS00_60", "CZCE_SA00_60", "CZCE_SH00_60", "CZCE_SF00_60", "CZCE_SM00_60",
    "CZCE_SR00_60", "CZCE_TA00_60", "CZCE_UR00_60", "CZCE_WH00_60", "CZCE_ZC00_60",
    "DCE_a00_60", "DCE_b00_60", "DCE_bb00_60", "DCE_bz00_60", "DCE_c00_60",
    "DCE_cs00_60", "DCE_eb00_60", "DCE_eg00_60", "DCE_fb00_60", "DCE_j00_60",
    "DCE_jd00_60", "DCE_jm00_60", "DCE_l00_60", "DCE_lg00_60", "DCE_lh00_60",
    "DCE_m00_60", "DCE_p00_60", "DCE_pg00_60", "DCE_pp00_60", "DCE_rr00_60",
    "DCE_v00_60", "DCE_y00_60", "GFEX_lc00_60", "GFEX_ps00_60", "GFEX_si00_60",
    "INE_bc00_60", "INE_ec00_60", "INE_lu00_60", "INE_nr00_60", "INE_sc00_60",
    "SHFE_ad00_60", "SHFE_ag00_60", "SHFE_al00_60", "SHFE_ao00_60", "SHFE_br00_60",
    "SHFE_bu00_60", "SHFE_cu00_60", "SHFE_fu00_60", "SHFE_hc00_60", "SHFE_ni00_60",
    "SHFE_op00_60", "SHFE_pb00_60", "SHFE_rb00_60", "SHFE_ru00_60", "SHFE_sn00_60",
    "SHFE_sp00_60", "SHFE_ss00_60", "SHFE_wr00_60", "SHFE_zn00_60"
]
# 跨周期数据合约池(日线数据)
extra_data_symbol = [
    "CZCE_AP00_86400", "CZCE_CF00_86400", "CZCE_CJ00_86400", "CZCE_CY00_86400", "CZCE_FG00_86400",
    "CZCE_LR00_86400", "CZCE_MA00_86400", "CZCE_OI00_86400", "CZCE_PF00_86400", "CZCE_PK00_86400",
    "CZCE_PL00_86400", "CZCE_PM00_86400", "CZCE_PR00_86400", "CZCE_RI00_86400", "CZCE_RM00_86400",
    "CZCE_RS00_86400", "CZCE_SA00_86400", "CZCE_SH00_86400", "CZCE_SF00_86400", "CZCE_SM00_86400",
    "CZCE_SR00_86400", "CZCE_TA00_86400", "CZCE_UR00_86400", "CZCE_WH00_86400", "CZCE_ZC00_86400",
    "DCE_a00_86400", "DCE_b00_86400", "DCE_bb00_86400", "DCE_bz00_86400", "DCE_c00_86400",
    "DCE_cs00_86400", "DCE_eb00_86400", "DCE_eg00_86400", "DCE_fb00_86400", "DCE_j00_86400",
    "DCE_jd00_86400", "DCE_jm00_86400", "DCE_l00_86400", "DCE_lg00_86400", "DCE_lh00_86400",
    "DCE_m00_86400", "DCE_p00_86400", "DCE_pg00_86400", "DCE_pp00_86400", "DCE_rr00_86400",
    "DCE_v00_86400", "DCE_y00_86400", "GFEX_lc00_86400", "GFEX_ps00_86400", "GFEX_si00_86400",
    "INE_bc00_86400", "INE_ec00_86400", "INE_lu00_86400", "INE_nr00_86400", "INE_sc00_86400",
    "SHFE_ad00_86400", "SHFE_ag00_86400", "SHFE_al00_86400", "SHFE_ao00_86400", "SHFE_br00_86400",
    "SHFE_bu00_86400", "SHFE_cu00_86400", "SHFE_fu00_86400"
]

# 加载数据
print('加载前', datetime.datetime.now().strftime("%H:%M:%S"))
base_data, all_trade_data = data.get_data(
    base_data_symbol,
    trade_data_symbol + extra_data_symbol,
    "20250101",
    "20251030"
)

# 分割数据
trade_data = {k: v for k, v in all_trade_data.items() if k in trade_data_symbol}
extra_data = {k: v for k, v in all_trade_data.items() if k in extra_data_symbol}
print('加载后', datetime.datetime.now().strftime("%H:%M:%S"))

# 回测参数配置
initial_cash = 10000 * 10000.0  # 初始资金: 1000万
fee_rate = 3.0 / 10000  # 手续费率: 万分之三

# 加载合约参数
with open("setting/volume_multiple.pkl", "rb") as f:
    loaded_volume = pickle.load(f)
with open("setting/margin_ratio.pkl", "rb") as f:
    loaded_margin = pickle.load(f)

deposit_rates = loaded_margin  # 保证金比例
multipliers = loaded_volume  # 合约乘数

print('向量化化开始', datetime.datetime.now().strftime("%H:%M:%S"))
# 交易数据计算
for symbol, data in trade_data.items():
    trade_data[symbol] = data.lazy().with_columns([
        pl.col("close").rolling_mean(5).alias("ma5"),
        pl.col("close").rolling_mean(10).alias("ma10"),
        pl.col("close").rolling_mean(5).shift(1).alias("prev_ma5"),
        pl.col("close").rolling_mean(10).shift(1).alias("prev_ma10"),
    ]).collect()

# 处理额外数据
for symbol, data in extra_data.items():
    extra_data[symbol] = data.lazy().with_columns([
        pl.col("close").rolling_mean(5).alias("ma5"),
        pl.col("close").rolling_mean(10).alias("ma10"),
        pl.col("close").rolling_mean(5).shift(1).alias("prev_ma5"),
        pl.col("close").rolling_mean(10).shift(1).alias("prev_ma10"),
    ]).collect()
print('向量化化完成', datetime.datetime.now().strftime("%H:%M:%S"))
class Strategy(PythonStrategy):
    def __init__(self, backtest_instance, trade_data, extra_data):

        super().__init__(backtest_instance, trade_data, extra_data)

        self.log_n = 0  # 日志计数器

    def on_bar(self, timestamp):
        """K线回调函数 - 每个时间戳调用一次"""
        if self.log_n % 10000 == 0:
            print('onbar前:{}, 时间:{}'.format(self.backtest_instance.today,datetime.datetime.now().strftime("%H:%M:%S")))
        # 遍历交易合约池,获取数据并计算信号
        for symbol in self.trade_data.keys():
            recent_data = self.get_recent_data_bisect(symbol, timestamp, lookback=11, require_exact_match=True)
            if recent_data is not None:
                self.calculate_signal(extract_symbol_base(symbol), recent_data, timestamp)
        if self.log_n % 10000 == 0:
            print('onbar后:{}, 时间:{}'.format(self.backtest_instance.today,datetime.datetime.now().strftime("%H:%M:%S")))
        self.log_n += 1

    def calculate_signal(self, symbol, recent_data, timestamp):
        # 直接从预计算的数据中读取指标值(性能优化)
        current_ma5 = recent_data["ma5"][-1]  # 当前5周期均线
        current_ma10 = recent_data["ma10"][-1]  # 当前10周期均线
        prev_ma5 = recent_data["prev_ma5"][-1]  # 前一个5周期均线
        prev_ma10 = recent_data["prev_ma10"][-1]  # 前一个10周期均线
        # 获取跨周期日线数据
        df_daily = self.get_recent_data_bisect(symbol + "_86400", timestamp, 2, False)
        if df_daily is None:
            return None
        if df_daily["ma5"][-1] is None:
            return None
        
        # 获取当前持仓
        current_holding = self.backtest_instance.long_position.get(symbol, 0)
        cond = df_daily["close"][-1] > df_daily["ma5"][-1]
        # 金叉信号开多
        if (prev_ma5 <= prev_ma10 and current_ma5 > current_ma10
                and current_holding == 0 and cond):
            self.backtest_instance.buy(symbol, 2, recent_data["close"][-1])

        # 死叉信号平仓
        elif (prev_ma5 >= prev_ma10 and current_ma5 < current_ma10
              and current_holding > 0):
            self.backtest_instance.sell(symbol, current_holding, recent_data["close"][-1])

def main():
    """
    主函数 - 配置并运行回测
    """
    # 初始化回测引擎
    my_backtest = backtest.BackTest(initial_cash, fee_rate, deposit_rates, multipliers, trade_data)

    # 创建策略实例 - 使用子类Strategy
    strategy = Strategy(my_backtest, trade_data, extra_data)

    # 准备回测时间序列
    datetime_strings = base_data.get_column("datetime").dt.strftime("%Y-%m-%d %H:%M:%S")

    # 启动回测
    backtest.start_backtest(strategy,my_backtest,datetime_strings,base_data.get_column("timestamps"),strategy_name)

if __name__ == "__main__":
    main()